import time
import pytest
from unittest.mock import MagicMock, patch
from pytest import MonkeyPatch

from motor.utils.data_builder import build_pod_ranktable, build_endpoints
from motor.resources.instance import Instance, ParallelConfig
from motor.resources.http_msg_spec import RegisterMsg, ReregisterMsg
from motor.controller.core.observer import ObserverEvent
from motor.controller.core.instance_assembler import InstanceAssembler


class TestInstanceAssembler:
    """Test InstanceAssembler with optimized performance"""
    
    def setup_method(self) -> None:
        """Setup for each test method"""
        self._cleanup_singleton()
        
        # save original config
        from motor.config.controller import controller_config
        self._original_config = {
            'instance_assemble_timeout': controller_config.instance_assemble_timeout,
            'send_cmd_retry_times': controller_config.send_cmd_retry_times,
            'max_link_number': controller_config.max_link_number,
            'instance_assembler_check_internal': controller_config.instance_assembler_check_internal,
            'instance_assembler_cmd_send_internal': controller_config.instance_assembler_cmd_send_internal,
        }
        
        self.dp = 4
        self.tp = 2
        self.role = "prefill"
        self.pod_ip1 = "127.0.0.1"
        self.pod_ip2 = "127.0.0.2"
        self.parallel_config = ParallelConfig(dp=self.dp, tp=self.tp)

    def _cleanup_singleton(self):
        """Clean up singleton instances"""
        if hasattr(InstanceAssembler, '_instances'):
            for assembler in list(InstanceAssembler._instances.values()):
                try:
                    assembler.stop()
                except Exception:
                    pass
            InstanceAssembler._instances.clear()
        time.sleep(0.01)

    def teardown_method(self) -> None:
        """Cleanup after each test"""
        # restore original config
        from motor.config.controller import controller_config
        if hasattr(self, '_original_config'):
            for key, value in self._original_config.items():
                setattr(controller_config, key, value)
        
        self._cleanup_singleton()

    @pytest.fixture
    def instance_assembler(self):
        """Setup mock assembler with optimized threading"""
        with patch('threading.Thread') as mock_thread_class:
            mock_thread = MagicMock()
            mock_thread_class.return_value = mock_thread
            assembler = InstanceAssembler()
            yield assembler
            try:
                assembler.stop()
            except Exception:
                pass

    @pytest.fixture
    def setup_http_server(self):
        """Mock HTTP server to avoid websockets warnings"""
        with patch('motor.controller.core.instance_assembler.SafeHTTPSClient') as mock_client_class:
            mock_client = MagicMock()
            mock_client.post.return_value = {'result': 'success'}
            mock_client_class.return_value = mock_client
            yield mock_client

    def test_register_succeed(self, instance_assembler: InstanceAssembler) -> None:
        """Test successful instance registration"""
        job_name = "testRegisterSucceed"

        result = instance_assembler.register(
            msg=RegisterMsg(
                job_name=job_name,
                model_name="test_model",
                role=self.role,
                pod_ip=self.pod_ip1,
                host_ip=self.pod_ip1,
                bussiness_port=["8080", "8084"],
                mgmt_port="9090",
                parallel_config=self.parallel_config,
                ranktable=build_pod_ranktable(pod_ip=self.pod_ip1, pod_device_num=2*self.tp)
            )
        )
        assert result == 0
        assert len(instance_assembler.instances) == 1

        result = instance_assembler.register(
            msg=RegisterMsg(
                job_name=job_name,
                model_name="test_model",
                role=self.role,
                pod_ip=self.pod_ip2,
                host_ip=self.pod_ip2,
                bussiness_port=["8080", "8084"],
                mgmt_port="9090",
                parallel_config=self.parallel_config,
                ranktable=build_pod_ranktable(
                    pod_ip=self.pod_ip2, 
                    pod_device_num=2 * self.tp, 
                    rank_offset=2 * self.tp,
                )
            )
        )
        assert result == 0

        instance = instance_assembler.instances[job_name]
        instance_assembler._assemble_instance(instance)
        
        assert len(instance_assembler.starting_instances) == 1
        assert len(instance_assembler.instances) == 0

    def test_register_timeout(self, instance_assembler: InstanceAssembler, monkeypatch: MonkeyPatch) -> None:
        """Test instance registration timeout"""
        monkeypatch.setattr('motor.config.controller.controller_config.instance_assemble_timeout', 0.05)
        
        job_name = "testRegisterTimeout"

        result = instance_assembler.register(
            msg=RegisterMsg(
                job_name=job_name,
                model_name="test_model",
                role=self.role,
                pod_ip=self.pod_ip1,
                host_ip=self.pod_ip1,
                bussiness_port=["8080"],
                mgmt_port="9090",
                parallel_config=self.parallel_config,
                ranktable=build_pod_ranktable(pod_ip=self.pod_ip1, pod_device_num=2*self.tp)
            )
        )
        assert result == 0
        assert len(instance_assembler.instances) == 1

        time.sleep(0.06)

        if job_name in instance_assembler.instances:
            instance_assembler._assemble_instance(instance_assembler.instances[job_name])
        
        assert len(instance_assembler.instances) == 0

    def test_reregister_succeed(self, instance_assembler: InstanceAssembler) -> None:
        """Test successful instance re-registration"""
        job_name = "testReregister"

        ep1 = build_endpoints(
            msg=RegisterMsg(
                job_name=job_name,
                model_name="test_model",
                role=self.role,
                pod_ip=self.pod_ip1,
                host_ip=self.pod_ip1,
                bussiness_port=["8080", "8084"],
                mgmt_port="9090",
                parallel_config=self.parallel_config,
                ranktable=build_pod_ranktable(pod_ip=self.pod_ip1, pod_device_num=2*self.tp)
            ),
        )
        ep2 = build_endpoints(
            msg=RegisterMsg(
                job_name=job_name,
                model_name="test_model",
                role=self.role,
                pod_ip=self.pod_ip2,
                host_ip=self.pod_ip2,
                bussiness_port=["8080", "8084"],
                mgmt_port="9090",
                parallel_config=self.parallel_config,
                ranktable=build_pod_ranktable(
                    pod_ip=self.pod_ip2, 
                    pod_device_num=2 * self.tp, 
                    rank_offset=2 * self.tp,
                )
            ),
            id_offset=self.tp
        )
        
        result = instance_assembler.reregister(
            msg=ReregisterMsg(
                job_name=job_name,
                model_name="test_model",
                instance_id=0,
                role=self.role,
                pod_ip=self.pod_ip1,
                host_ip=self.pod_ip1,
                mgmt_port="9090",
                parallel_config=self.parallel_config,
                endpoints=[endpoint for endpoint in ep1.values()]
            )
        )
        assert result == 0
        assert len(instance_assembler.instances) == 1

        result = instance_assembler.reregister(
            msg=ReregisterMsg(
                job_name=job_name,
                model_name="test_model",
                instance_id=0,
                role=self.role,
                pod_ip=self.pod_ip2,
                host_ip=self.pod_ip2,
                mgmt_port="9090",
                parallel_config=self.parallel_config,
                endpoints=[endpoint for endpoint in ep2.values()]
            )
        )

        assert result == 0
        assert instance_assembler.ins_id_cnt == 1

        instance_assembler._assemble_instance(instance_assembler.instances[job_name])
        assert len(instance_assembler.instances) == 0

    def test_send_start_cmd(self, instance_assembler: InstanceAssembler, setup_http_server) -> None:
        """Test start command sending"""
        job_name = "testSendStartCommand"

        result = instance_assembler.register(
            msg=RegisterMsg(
                job_name=job_name,
                model_name="test_model",
                role=self.role,
                pod_ip=self.pod_ip1,
                host_ip=self.pod_ip1,
                bussiness_port=["8080", "8084"],
                mgmt_port="9090",
                parallel_config=self.parallel_config,
                ranktable=build_pod_ranktable(pod_ip=self.pod_ip1, pod_device_num=2*self.tp)
            )
        )
        assert result == 0

        instance = instance_assembler.instances[job_name]
        instance_assembler.starting_instances.add(job_name)
        
        success = instance_assembler._send_start_command(instance)
        assert success == True

    def test_send_start_cmd_fail_retry(self, instance_assembler: InstanceAssembler) -> None:
        """Test start command failure retry"""
        job_name = "testSendStartCommandFailRetry"

        result = instance_assembler.register(
            msg=RegisterMsg(
                job_name=job_name,
                model_name="test_model",
                role=self.role,
                pod_ip=self.pod_ip1,
                host_ip=self.pod_ip1,
                bussiness_port=["8080", "8084"],
                mgmt_port="9091",
                parallel_config=self.parallel_config,
                ranktable=build_pod_ranktable(pod_ip=self.pod_ip1, pod_device_num=2*self.tp)
            )
        )
        assert result == 0

        instance_assembler.assembled_instances.put((instance_assembler.instances[job_name], 0))
        instance_assembler.assembled_instances.put(None)
        assert instance_assembler.assembled_instances.qsize() == 2

        def mock_stop_sleep(seconds):
            if instance_assembler.assembled_instances.qsize() > 0:
                raise StopIteration

        with patch('motor.controller.core.instance_assembler.time') as mock_time:
            mock_time.sleep.side_effect = mock_stop_sleep
            try:
                instance_assembler._start_commmand_sender()
            except StopIteration:
                pass

            assert instance_assembler.assembled_instances.qsize() == 2

    def test_send_start_cmd_fail_abort(self, instance_assembler: InstanceAssembler, monkeypatch: MonkeyPatch) -> None:
        """Test start command failure abort"""
        from motor.config.controller import controller_config
        monkeypatch.setattr(controller_config, 'send_cmd_retry_times', 1)
        job_name = "testSendStartCommandFailAbort"

        result = instance_assembler.register(
            msg=RegisterMsg(
                job_name=job_name,
                model_name="test_model",
                role=self.role,
                pod_ip=self.pod_ip1,
                host_ip=self.pod_ip1,
                bussiness_port=["8080", "8084"],
                mgmt_port="9091",
                parallel_config=self.parallel_config,
                ranktable=build_pod_ranktable(pod_ip=self.pod_ip1, pod_device_num=2*self.tp)
            )
        )
        assert result == 0

        instance_assembler.assembled_instances.put((instance_assembler.instances[job_name], 0))
        instance_assembler.assembled_instances.put(None)
        assert instance_assembler.assembled_instances.qsize() == 2

        def mock_stop_sleep(seconds):
            if instance_assembler.assembled_instances.qsize() > 0:
                raise StopIteration

        with patch('motor.controller.core.instance_assembler.time') as mock_time:
            mock_time.sleep.side_effect = mock_stop_sleep
            try:
                instance_assembler._start_commmand_sender()
            except StopIteration:
                pass

            assert instance_assembler.assembled_instances.qsize() == 1

    def test_alloc_ins_group_success(self, instance_assembler: InstanceAssembler, monkeypatch: MonkeyPatch) -> None:
        """Test successful instance group allocation"""
        from motor.config.controller import controller_config
        monkeypatch.setattr(controller_config, 'max_link_number', 8)
        
        instance1 = Instance(
            job_name="testAllocInsGroup1",
            model_name="test_model",
            id=1,
            role=self.role,
            parallel_config=self.parallel_config
        )

        instance2 = Instance(
            job_name="testAllocInsGroup2",
            model_name="test_model",
            id=2,
            role=self.role,
            parallel_config=ParallelConfig(dp=self.dp, tp=self.tp / 2)
        )
        instance3 = Instance(
            job_name="testAllocInsGroup3",
            model_name="test_model",
            id=3,
            role=self.role,
            parallel_config=ParallelConfig(dp=self.dp, tp=self.tp / 2)
        )

        instance_assembler._assign_ins_group(instance1)
        instance_assembler._assign_ins_group(instance1)

        instance_assembler._assign_ins_group(instance2)
        instance_assembler._assign_ins_group(instance3)

        group_info = instance_assembler.instances_group

        assert len(group_info) == 2
        assert group_info[0].current_group_member == self.dp * self.tp
        assert group_info[1].current_group_member == self.dp * self.tp

    def test_alloc_ins_group_fail(self, instance_assembler: InstanceAssembler, monkeypatch: MonkeyPatch) -> None:
        """Test failed instance group allocation"""
        max_link_number = 2
        from motor.config.controller import controller_config
        monkeypatch.setattr(controller_config, 'max_link_number', max_link_number)

        instance = Instance(
            job_name="testAllocInsGroupFail",
            model_name="test_model",
            id=1,
            role=self.role,
            parallel_config=self.parallel_config
        )

        try:
            instance_assembler._assign_ins_group(instance)
        except Exception as e:
            error_msg = (
                f"Instance {instance.job_name}(id:{instance.id}) allocate ins group failed, "
                f"max link number is {max_link_number}, but need {instance.parallel_config.world_size}."
            )
            assert str(e) == error_msg
            return

        assert False, "Expected exception was not raised."

    def test_ins_group_metadata_update(self, instance_assembler: InstanceAssembler, monkeypatch: MonkeyPatch) -> None:
        """Test instance group metadata update"""
        from motor.config.controller import controller_config
        monkeypatch.setattr(controller_config, 'max_link_number', 8)
        
        instance1 = Instance(
            job_name="testAllocInsGroup1",
            model_name="test_model",
            id=1,
            role="prefill",
            parallel_config=ParallelConfig(dp=self.dp, tp=self.tp / 2)
        )
        instance2 = Instance(
            job_name="testInsGroupMetadataUpdate",
            model_name="test_model",
            id=2,
            role="decode",
            parallel_config=ParallelConfig(dp=self.dp, tp=self.tp / 2)
        )
        
        instance_assembler._assign_ins_group(instance1)
        instance_assembler._assign_ins_group(instance2)
        
        assert len(instance_assembler.instances_group) == 1
        assert instance_assembler.instances_group[0].current_group_member == self.dp * self.tp
        assert instance1.group_id == 0 and instance2.group_id == 0

        instance_assembler.update(instance1, ObserverEvent.INSTANCE_REMOVED)
        instance_assembler.update(instance2, ObserverEvent.INSTANCE_REMOVED)
        assert instance_assembler.instances_group[0].current_group_member == 0

        instance_assembler.update(instance2, ObserverEvent.INSTANCE_ADDED)
        assert instance_assembler.instances_group[0].current_group_member == self.dp * self.tp / 2
        assert instance2.group_id == 0

    def test_ins_group_metadata_update_exception_handle(self, instance_assembler: InstanceAssembler, monkeypatch: MonkeyPatch) -> None:
        """Test instance group metadata update exception handling"""
        from motor.config.controller import controller_config
        monkeypatch.setattr(controller_config, 'max_link_number', 8)
        
        instance1 = Instance(
            job_name="testInsGroupMetadataUpdateExceptionHandle",
            model_name="test_model",
            id=1,
            role="prefill",
            parallel_config=ParallelConfig(dp=self.dp, tp=self.tp)
        )
        
        instance_assembler._assign_ins_group(instance1)
        assert len(instance_assembler.instances_group) == 1
        assert instance_assembler.instances_group[0].current_group_member == self.dp * self.tp
        assert instance1.group_id == 0

        instance_assembler.update(instance1, ObserverEvent.INSTANCE_REMOVED)
        assert instance_assembler.instances_group[0].current_group_member == 0
        instance_assembler.update(instance1, ObserverEvent.INSTANCE_REMOVED)
        assert instance_assembler.instances_group[0].current_group_member == 0

        try:
            instance_assembler.update(instance1, {})
        except Exception as e:
            error_msg = f"Unknown event type: {{}}"
            assert str(e) == error_msg
            return

        assert False, "Expected exception was not raised."

    def test_exception_handle(self, instance_assembler: InstanceAssembler):
        """Test exception handling"""
        try:
            instance_assembler.register({})
        except Exception as e:
            error_msg = (f"Invalid msg provided to register. "
                         f"expect RegisterMsg, got {type({})}")
            assert str(e) == error_msg
            return
        try:
            instance_assembler.reregister({})
        except Exception as e:
            error_msg = (f"Invalid msg provided to reregister. "
                         f"expect ReregisterMsg, got {type({})}")
            assert str(e) == error_msg
            return
        
        try:
            instance_assembler._assemble_instance({})
        except Exception as e:
            error_msg = (f"Invalid instance provided to assemble. "
                         f"expect Instance, got {type({})}")
            assert str(e) == error_msg
            return 
        
        assert False, "Expected exception was not raised."

    def test_stop(self, monkeypatch: MonkeyPatch) -> None:
        """Test stop method"""
        from motor.config.controller import controller_config
        monkeypatch.setattr(controller_config, 'instance_assembler_check_internal', 0.01)
        monkeypatch.setattr(controller_config, 'instance_assembler_cmd_send_internal', 0.01)
        
        instance_assembler = InstanceAssembler()
        instance_assembler.stop()

        assert instance_assembler.stop_event.is_set() == True

    def test_performance(self, instance_assembler: InstanceAssembler, monkeypatch: MonkeyPatch) -> None:
        """Test performance under load"""
        monkeypatch.setattr('motor.config.controller.controller_config.max_link_number', 100)
        monkeypatch.setattr('motor.config.controller.controller_config.instance_assemble_timeout', 10)
        
        num_instances = 10
        start_time = time.time()
        
        for i in range(num_instances):
            result = instance_assembler.register(
                RegisterMsg(
                    job_name=f"perf_test_{i}",
                    model_name="test_model",
                    role=self.role,
                    pod_ip=f"127.0.0.{i % 5}",
                    host_ip=f"127.0.0.{i % 5}",
                    bussiness_port=["8080", "8084"],
                    mgmt_port="9090",
                    parallel_config=self.parallel_config,
                    ranktable=build_pod_ranktable(
                        pod_ip=f"127.0.0.{i % 5}", 
                        pod_device_num=2*self.tp
                    )
                )
            )
            assert result == 0
        
        end_time = time.time()
        registration_time = end_time - start_time
        
        assert len(instance_assembler.instances) == num_instances
        assert registration_time < 5.0

    def test_concurrent_registration(self, instance_assembler: InstanceAssembler) -> None:
        """Test concurrent registration"""
        job_name = "testConcurrentRegistration"
        num_instances = 3
        
        results = []
        for i in range(num_instances):
            result = instance_assembler.register(
                RegisterMsg(
                    job_name=f"{job_name}_{i}",
                    model_name="test_model",
                    role=self.role,
                    pod_ip=f"127.0.0.{i}",
                    host_ip=f"127.0.0.{i}",
                    bussiness_port=["8080", "8084"],
                    mgmt_port="9090",
                    parallel_config=self.parallel_config,
                    ranktable=build_pod_ranktable(
                        pod_ip=f"127.0.0.{i}", 
                        pod_device_num=2*self.tp
                    )
                )
            )
            results.append(result)
        
        assert all(result == 0 for result in results)
        assert len(instance_assembler.instances) == num_instances